﻿using System;
using System.Collections.Generic;
using System.Drawing;
using Emgu.CV;
using Emgu.CV.CvEnum;
using Emgu.CV.Structure;

namespace ImageAlignment
{
  class InverseCompositionalAlgorithm
  {
    public static void Run(Image<Gray, Byte> templateImage, Rectangle templateRect, Image<Gray, Byte> image)
    {
      // Some constants for iterative minimization process.
      const float eps = 1E-5f; // Threshold value for termination criteria.
      const int maxIter = 100;  // Maximum iteration count.

      // Here we will store internally used images.
      // Gradient of I in X direction.
      var gradIx = new Image<Gray, short>(image.Width, image.Height);
      // Gradient of I in Y direction.
      var gradIy = new Image<Gray, short>(image.Width, image.Height);
      // Steepest descent images.
      var stDesc0 = new Image<Gray, float>(image.Width, image.Height);
      var stDesc1 = new Image<Gray, float>(image.Width, image.Height);
      var stDesc2 = new Image<Gray, float>(image.Width, image.Height);
      var stDesc = new Image<Gray, float>(new[] { stDesc0, stDesc1, stDesc2 });

      //TODO:Make 3 channels instead of 1.

      // Create matrices.
      var W = new Matrix<float>(3, 3);//Current value of warp W(x,p)
      var dW = new Matrix<float>(3, 3);//Warp update.
      var idW = new Matrix<float>(3, 3);//Inverse of warp update.
      var X = new Matrix<float>(3, 1);//Point in coordinate frame of T
      var Z = new Matrix<float>(3, 1);//Point in coordinate frame of I

      var H = new Matrix<float>(3, 3);  // Hessian
      var iH = new Matrix<float>(3, 3); // Inverse of Hessian
      var b = new Matrix<float>(3, 1);  // Vector in the right side of the system of linear equations.
      var delta_p = new Matrix<float>(3, 1); // Parameter update value.

      var imageArray = AuxFunction.GetArray(image);

      var start_time = DateTime.Now;

      CvInvoke.cvSobel(templateImage.Ptr, gradIx.Ptr, 1, 0, 3);
      CvInvoke.cvConvertScale(gradIx.Ptr, gradIx.Ptr, 0.125, 0); // Normalize
      CvInvoke.cvSobel(templateImage.Ptr, gradIy.Ptr, 1, 0, 3);
      CvInvoke.cvConvertScale(gradIy.Ptr, gradIy.Ptr, 0.125, 0); // Normalize

      // Compute steepest descent images and Hessian
      AuxFunction.SetAll(H, 0);

      int u, v;// (u,v) - pixel coordinates in the coordinate frame of T.
      float u2, v2;// (u2,v2) - pixel coordinates in the coordinate frame of I

      // Walk through pixels in the template T.
      int i, j;
      for (i = 0; i < templateRect.Width; i++)
      {
        u = i + templateRect.X;

        for (j = 0; j < templateRect.Height; j++)
        {
          v = j + templateRect.Y;

          //Evaluate gradient of T ?????
          short Tx = (short)gradIx[v, u].Intensity;
          short Ty = (short)gradIy[v, u].Intensity;

          // Calculate steepest descent image's element.
          //var stDescArr = float.Parse(stDesc[v,3*u].ToString());// an element of steepest descent image
          var stdesc = new float[3]; // an element of steepest descent image
          stdesc[0] = (float)(-v * Tx + u * Ty);
          stdesc[1] = (float)Tx;
          stdesc[2] = (float)Ty;

          //Add a term to Hessian
          int l, m;
          for (l = 0; l < 3; l++)
          {
            for (m = 0; m < 3; m++)
            {
              H[l, m] += stdesc[l] * stdesc[m];
            }
          }
        }
      }

      // Invert Hessian.
      double inv_res = CvInvoke.cvInvert(H.Ptr, iH.Ptr, INVERT_METHOD.CV_LU);
      if (inv_res == 0)
      {
        Console.WriteLine("Error: Hessian is singular.\n");
        return;
      }

      /*
       *   Iteration stage.
       */

      // Set warp with identity.
      AuxFunction.SetIdentity(W, 1);

      // Here we will store current value of mean error.
      float mean_error = 0;

      // Iterate
      int iter = 0; // number of current iteration
      while (iter < maxIter)
      {
        iter++;//Increment iteration counter

        mean_error = 0;// Set mean error value with zero

        int pixel_count = 0;// Count of processed pixels

        AuxFunction.SetAll(b, 0); //Set b matrix with zeroes

        //Walk through pixels in the template T
        for (i = 0; i < templateRect.Width; i++)
        {
          u = i + templateRect.X;
          for (j = 0; j < templateRect.Height; j++)
          {
            v = j + templateRect.Y;

            // Set vector X with pixel coordinates (u,v,1)
            X = AuxFunction.CreateVector(u, v, 1);

            // Warp Z=W*X
            Z = W.Mul(X);

            // Get coordinates of warped pixel in coordinate frame of I.
            u2 = Z[0, 0];
            v2 = Z[1, 0];

            // Get the nearest integer pixel coords (u2i;v2i).
            int u2i = (int)Math.Floor(u2);
            int v2i = (int)Math.Floor(v2);

            if (u2i >= 0 && u2i < image.Width && // check if pixel is inside I.
          v2i >= 0 && v2i < image.Height)
            {
              pixel_count++;

              // Calculate intensity of a transformed pixel with sub-pixel accuracy
              // using bilinear interpolation.
              float I2 = AuxFunction.Interpolate(imageArray, u2, v2);

              // Calculate image difference D = I(W(x,p))-T(x).
              float D = I2 - (float)templateImage[v, u].Intensity;

              //Update mean error value.
              mean_error += Math.Abs(D);


              //Evaluate gradient of T ?????
              short Tx = (short)gradIx[v, u].Intensity;
              short Ty = (short)gradIy[v, u].Intensity;

              // Calculate steepest descent image's element.
              //var stDescArr = float.Parse(stDesc[v,3*u].ToString());// an element of steepest descent image
              var stdesc = new float[3]; // an element of steepest descent image
              stdesc[0] = (float)(-v * Tx + u * Ty);
              stdesc[1] = (float)Tx;
              stdesc[2] = (float)Ty;

              //Add a term to b matrix
              b[0, 0] += stdesc[0] * D;
              b[1, 0] += stdesc[1] * D;
              b[2, 0] += stdesc[2] * D;
            }
          }
        }
        // Finally, calculate resulting mean error.
        if (pixel_count != 0)
          mean_error /= pixel_count;

        // Find parameter increment. 
        delta_p = iH.Mul(b);
        float delta_wz = delta_p[0, 0];
        float delta_tx = delta_p[1, 0];
        float delta_ty = delta_p[2, 0];

        AuxFunction.InitWarp(dW, delta_wz, delta_tx, delta_ty);
        //Invert warp.
        inv_res = CvInvoke.cvInvert(dW.Ptr, idW.Ptr, INVERT_METHOD.CV_LU);
        if (inv_res == 0)
        {
          Console.WriteLine("Error:Warp matrix is singular.\n");
          return;
        }

        dW = idW.Mul(W);

        W = dW;

        //Print diagnostic infromation to screen.
        Console.WriteLine("iter={0} mean_error={1}\n", iter, mean_error);

        //Check termination criteria.
        if (Math.Abs(delta_wz) <= eps && Math.Abs(delta_tx) <= eps && Math.Abs(delta_ty) <= eps) break;
      }

      // Get current time and obtain total time of calculation.
      var finish_time = DateTime.Now;
      double total_time = (finish_time - start_time).TotalSeconds;

      // Print summary.
      Console.WriteLine("===============================================\n");
      Console.WriteLine("Algorithm: inverse compositional.\n");
      Console.WriteLine("Caclulation time: {0} sec.\n", total_time);
      Console.WriteLine("Iteration count: {0}\n", iter);
      Console.WriteLine("Epsilon: {0}\n", eps);
      Console.WriteLine("Resulting mean error: {0}\n", mean_error);
      Console.WriteLine("===============================================\n");

      // Show result of image alignment.
      AuxFunction.DrawWarpedRectangle(image, templateRect, W);
      CvInvoke.cvSetImageROI(templateImage, templateRect);
      CvInvoke.cvShowImage("template", templateImage);
      CvInvoke.cvShowImage("image", image);
      CvInvoke.cvResetImageROI(templateImage);
    }
  }
}